import numpy as np
import matplotlib.pyplot as plt

def elliptical_hole_plate_stress_distribution():
    # 椭圆挖空无限大板的解析解公式
    def elliptical_hole_plate_stress_analytical(x, y, a, b, P):
        dx = abs(x)-a
        dy = y
        rho = np.sqrt(dx**2 + dy**2)  # 计算距离
        theta = np.arctan2(dy, dx)     # 计算角度

        stress_with_hole = P*(2*rho/a)**-0.5*np.cos(theta/2)*(1+np.sin(theta/2)*np.sin(theta*1.5))
        return stress_with_hole

    # 定义椭圆孔板和加载参数
    a = 0.1  # 椭圆长半轴
    b = 0.05  # 椭圆短半轴
    P = 1.0  # 受力

    # 绘制应力分布图
    x_values = np.linspace(-2 * a, 2 * a, 2000)
    y_values = np.linspace(-2 * b, 2 * b, 2000)
    X, Y = np.meshgrid(x_values, y_values)
    stress = elliptical_hole_plate_stress_analytical(X, Y, a, b, P)

    # 将不在椭圆内的点的 stress 设为 NaN
    ellipse_mask = (X**2 / a**2) + (Y**2 / b**2) < 1
    stress[ellipse_mask] = np.nan

    plt.figure(figsize=(12, 5))

    # 左侧：解析解公式和椭圆图形
    plt.subplot(1, 2, 1)
    plt.text(0.5, 0.9, r'$\sigma_y = \frac{4P}{\pi ab}\left(\frac{a^2 - b^2}{a^2 + b^2}\right)\left(\frac{x^2}{a^2} + \frac{y^2}{b^2}\right)$', fontsize=14, ha='center')
    theta = np.linspace(0, 2*np.pi, 100)
    
    plt.axis('equal')  # 保持坐标轴比例一致
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title('Analytical Solution and Ellipse')
    plt.axis('equal')

    # 右侧：应力分布图
    upper_color = 6
    lower_color = 0
    plt.subplot(1, 2, 2)
    plt.contourf(X, Y, stress, cmap='viridis', levels=20, vmin=lower_color, vmax=upper_color)
    plt.plot(a*np.cos(theta), b*np.sin(theta), color='red', linestyle='--', label='Ellipse')
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title('Stress Distribution')
    plt.grid(True)
    plt.axis('equal')

    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    # 测试函数
    elliptical_hole_plate_stress_distribution()
